# install packages from CRAN
packages <- c("tidyverse", "haven", "lubridate", "tigris", "fixest", 
              "tidycensus", "usmap", "ggpattern", "cowplot", "ggpubr")

install.packages(setdiff(packages, rownames(installed.packages())))  

# Create two folders, one for figures, one for tables
if(length(dir(pattern = "figures"))==0){
  dir.create("figures")
}else{
  print("Folder already created.")
}


### Figure S21 ###

# Pre-trend Responses, Ever Vaccinated versus Never Vaccinated
library(lubridate)

data <- read_dta("data_gallup.dta")

# Keep variables of interest
d <- data %>%
  select(Mostly_Isol, worn_mask, v_worry_ill, mostly_remote,
         pid, time_string, time, 
         dem, gop, indep_third, 
         vaccinated, WEIGHT,  ctyfip,
         ZLNpc_new_cases_7day, out_workforce, employed, income, 
         gop, dem, indep_third,
         Yc4_restrictionsongatherings, Yc6_stayathomerequirements) %>%
  mutate(dem_vac = dem * vaccinated,
         gop_vac = gop * vaccinated, 
         ind_vac = indep_third * vaccinated)

# Treated respondents if ever vaccinated
treated = d %>%
  group_by(pid) %>%
  mutate(treated = if_else(sum(vaccinated, na.rm = T) > 0, 1, 0)) %>%
  ungroup()

# combine treated and control respondents
pt = bind_rows(treated %>% filter(treated == 1) %>% filter(vaccinated == 0), # treated respondents before they got the vaccine
               treated %>% filter(treated == 0)
)


# before vaccination roll-out Feb 1 2021
pt = pt %>%
  filter(time_string <= "2021-02-01") %>%
  select(time_string, treated, 
         Mostly_Isol, worn_mask, v_worry_ill, mostly_remote) %>%
  pivot_longer(!time_string:treated, names_to = "dv", values_to = "value")

pt$dv = case_when(pt$dv == "Mostly_Isol" ~ "Mostly isolating",
                  pt$dv == "mostly_remote" ~ "Mostly working remote",
                  pt$dv == "v_worry_ill" ~ "Very worried over COVID",
                  pt$dv == "worn_mask" ~ "Worn mask")

ggplot(pt, aes(x = as.POSIXct(time_string), y = value, 
               color = factor(treated), 
               fill = factor(treated), 
               shape = factor(treated), 
               lintype = factor(treated))) +
  stat_smooth(linewidth = .8) +
  facet_wrap(~dv, ncol = 2, scale = "free_y") +
  stat_summary_bin(fun.data = mean_se, geom = "point", binwidth = 15*24*60*60, alpha = .3, size = 1.2) +
  theme_minimal() +
  coord_cartesian(ylim = c(0, 1), expand = T) +
  labs(x = "Date", y = "Average Share of Respondents\nAnswering Affirmatively", 
       shape = NULL, fill = NULL, color = NULL, linetype = NULL) +
  scale_color_manual(values = c("darkgreen", "darkred"),
                     labels = c("0" = "Never Vaccinated", "1" = "Ever Vaccinated")) +
  scale_fill_manual(values = c("darkgreen", "darkred"),
                    labels = c("0" = "Never Vaccinated", "1" = "Ever Vaccinated")) +
  scale_shape_manual(values = c(16, 17),
                     labels = c("0" = "Never Vaccinated", "1" = "Ever Vaccinated")) +
  scale_linetype_manual(values = c("solid", "dashed"),
                        labels = c("0" = "Never Vaccinated", "1" = "Ever Vaccinated")) +  
  theme(panel.grid = element_blank(), 
        panel.border = element_rect(fill = NA), 
        plot.title = element_text(hjust = .5, face = "bold", size = 9),
        strip.text = element_text(size = 8), 
        axis.text = element_text(size = 6),
        axis.title = element_text(size = 8),
        legend.text = element_text(size = 7),
        legend.position = "bottom",
        legend.key.height= unit(.3, 'cm'),
        legend.key.width= unit(.3, 'cm')) +
  scale_x_datetime(date_labels = "%b %y")

ggsave("figures/Figure S21.png", bg = "white", width = 5.5, height = 4)


### Figure S24 ###

rm(list = ls()[ls()!="lf"]) # do not remove log file

# Map of Gatherings Restrictions and Stay-at-Home Orders 

library(usmap)
library(ggpattern)
library(cowplot)
library(ggpubr)

# Create a data frame with state abbreviations and policy types
state_policies <- data.frame(
  state = state.abb,
  party = c("R11", "R11", "R00", "R10", "D11", "D11", "D10", "D11", "R11", "R01",
            "D11", "R11", "D11", "R11", "R10", "RXD10", "RXD11", "RXD11", "R11", "DXR10",
            "MA", "RXD11", "D11", "R00", "R11", "RXD11", "R10", "D10", "DXR11", "D11",
            "D10", "D11", "RXD11", "R00", "R11", "R11", "D11", "RXD11", "RI", "R11",
            "R10", "R11", "R11", "R10", "DXR11", "D11", "D11", "R11", "RXD11", "R10")
)



################################
############ MAP ###############
################################

m <- plot_usmap(data = state_policies, values = "party") +
  geom_sf_pattern(
    aes(pattern = party, fill = party, pattern_fill = party, pattern_color = party,
        pattern_density = party, pattern_spacing = party), 
    color = "gray70", pattern_alpha = 0.7
  ) +
  scale_pattern_manual(values = c(
    "R11" = "crosshatch", "R10" = "circle", "R00" = "none", "R01" = "stripe",
    "D11" = "crosshatch", "D10" = "circle",
    "RXD10" = "circle", "RXD11" = "crosshatch", "DXR10" = "circle", "DXR11" = "crosshatch", 
    "MA" = "circle", "RI" = "crosshatch"
  )) +
  scale_pattern_density_manual(values = c(
    "R11" = 0.02, "R10" = 0.1, "R00" = 0, "R01" = 0.1,
    "D11" = 0.02, "D10" = 0.1,
    "RXD10" = 0.1, "RXD11" = 0.02, "DXR10" = 0.1, "DXR11" = 0.02, 
    "MA" = 0.1, "RI" = 0.03
  )) +
  scale_pattern_spacing_manual(values = c(
    "R11" = 0.04, "R10" = 0.02, "R00" = 0, "R01" = 0.02,
    "D11" = 0.04, "D10" = 0.02,
    "RXD10" = 0.02, "RXD11" = 0.04, "DXR10" = 0.02, "DXR11" = 0.04, 
    "MA" = 0.02, "RI" = 0.04
  )) +
  scale_fill_manual(values = c(
    "D11" = alpha("#005496", 0.2), "R11" = alpha("red3", 0.2),
    "R10" = alpha("red3", 0.2), "R00" = alpha("red3", 0.2),
    "R01" = alpha("red3", 0.2), "D10" = alpha("#005496", 0.2),
    "RXD11" = alpha("red3", 0.2), "RXD10" = alpha("red3", 0.2),
    "DXR10" = alpha("#005496", 0.2), "DXR11" = alpha("#005496", 0.2),
    "MA" = alpha("#005496", 0.2), "RI" = alpha("#005496", 0.2)
  )) +
  scale_pattern_color_manual(values = c(
    "R11" = "red3", "R10" = "red3", "R00" = "red3", "R01" = "red3",
    "D11" = "blue", "D10" = "blue", "RXD10" = "blue", "RXD11" = "blue",
    "DXR10" = "red3", "DXR11" = "red3", "MA" = "red3", "RI" = "blue"
  )) +
  theme(
    legend.position = "none",  # This removes the legend
    text = element_text(size = 8)
  ) +
  guides(fill = "none", pattern = "none")


################################
########## LEGEND ##############
################################

df <- data.frame(
  trt = c("Democratic state", "Republican state", 
          "Both gathering restriction and stay-at-home order", 
          "Only gathering restriction", 
          "Only stay-at-home order"), 
  outcome = c(2.3, 1.9, 3.2, 1, 2)
)

# Convert to factor with specific order
df$trt <- factor(df$trt, levels = c("Democratic state", "Republican state", 
                                    "Both gathering restriction and stay-at-home order", 
                                    "Only gathering restriction", 
                                    "Only stay-at-home order"))

# Create plot just for the legend
legend_plot <- ggplot(df, aes(x = trt, y = outcome)) +
  geom_col_pattern(
    aes(fill = trt, pattern = trt),
    colour = 'white', 
    pattern_density = 0.1,
    pattern_spacing = 0.03,
    pattern_fill = 'black'
  ) +
  theme_bw() +
  scale_fill_manual(values = c('#77b7e6', '#ed7872', 'white', 'white', 'white')) +
  scale_pattern_manual(values = c('none', 'none', 'weave', 'circle', 'stripe')) +
  theme(
    legend.key.size = unit(1, 'cm'),
    legend.key = element_rect(colour = "black"),
    legend.title = element_blank()
  )

# Extract just the legend
legend_only <- get_legend(legend_plot)

# Combine
combined <- plot_grid(
  m, legend_only,
  ncol = 2,    
  rel_widths = c(2, 0.5)
)

ggsave("figures/Figure S24.png", bg = "white", plot = combined, width = 10, height = 6, dpi = 300)


### Figure S25 ###

rm(list = ls()[ls()!="lf"]) # do not remove log file

# Nominal versus Effective Samples for Table 2

library(fixest)
library(tidycensus)
library(tigris)

# Load data
d <- read_dta("data_gallup.dta")

# recode party
d <- d %>%
  mutate(party_rep = case_when(party == 1 ~ "D",
                               party == 2 ~ "I",
                               party == 3 ~ "R"))


# drop respondents with missing obs on key variables used for analysis
das_gatherings = d %>%
  select(approvestateresponse, party_rep, dem, gop, indep, ZLNpc_new_cases_7day, out_workforce, employed, male, age10, age_group4, somecol, ba, grad, AmerInd, Asian, Black, Hisp, Multiracial, live_w_children, income, stateid, state, time, gopgovernor, demgovernor, WEIGHT,
         demrestrictionsongatherings, indrestrictionsongatherings, Yc4_restrictionsongatherings) %>%
  drop_na()

das_saho = d %>%
  select(approvestateresponse, party_rep, dem, gop, indep, ZLNpc_new_cases_7day, out_workforce, employed, male, age10, age_group4, somecol, ba, grad, AmerInd, Asian, Black, Hisp, Multiracial, live_w_children, income, stateid, state, time, gopgovernor, demgovernor, WEIGHT,
         Yc6_stayathomerequirements) %>%
  drop_na()

# Estimate regression where policy treatment is outcome, regressed on battery of covariates and fixed effects as in main specification of Table 2

# Gatherings
as_gatherings = feols(Yc4_restrictionsongatherings ~ dem + indep +
                        ZLNpc_new_cases_7day 
                      + out_workforce + employed + male + age10 + age_group4  +  somecol + ba + grad + AmerInd + Asian + Black + Hisp + Multiracial + live_w_children + income | stateid + time,
                      weight = ~ WEIGHT, 
                      cluster = ~ stateid, 
                      data = das_gatherings)

as_saho  = feols(Yc6_stayathomerequirements ~ dem + indep +
                   ZLNpc_new_cases_7day 
                 + out_workforce + employed + male + age10 + age_group4  +  somecol + ba + grad + AmerInd + Asian + Black + Hisp + Multiracial + live_w_children + income | stateid + time,
                 weight = ~ WEIGHT, 
                 cluster = ~ stateid, 
                 data = das_saho)

# Dem Governors
as_gatherings_D = feols(Yc4_restrictionsongatherings ~ dem + indep +
                          ZLNpc_new_cases_7day 
                        + out_workforce + employed + male + age10 + age_group4  +  somecol + ba + grad + AmerInd + Asian + Black + Hisp + Multiracial + live_w_children + income | stateid + time,
                        subset = ~ demgovernor == 1, 
                        weight = ~ WEIGHT, 
                        cluster = ~ stateid, 
                        data = das_gatherings)

as_saho_D  = feols(Yc6_stayathomerequirements ~ dem + indep +
                     ZLNpc_new_cases_7day 
                   + out_workforce + employed + male + age10 + age_group4  +  somecol + ba + grad + AmerInd + Asian + Black + Hisp + Multiracial + live_w_children + income | stateid + time,
                   subset = ~ demgovernor == 1, 
                   weight = ~ WEIGHT, 
                   cluster = ~ stateid, 
                   data = das_saho)

# GOP Governors
as_gatherings_R = feols(Yc4_restrictionsongatherings ~ dem + indep +
                          ZLNpc_new_cases_7day 
                        + out_workforce + employed + male + age10 + age_group4  +  somecol + ba + grad + AmerInd + Asian + Black + Hisp + Multiracial + live_w_children + income | stateid + time,
                        subset = ~ gopgovernor == 1, 
                        weight = ~ WEIGHT, 
                        cluster = ~ stateid, 
                        data = das_gatherings)

as_saho_R  = feols(Yc6_stayathomerequirements ~ dem + indep +
                     ZLNpc_new_cases_7day 
                   + out_workforce + employed + male + age10 + age_group4  +  somecol + ba + grad + AmerInd + Asian + Black + Hisp + Multiracial + live_w_children + income | stateid + time,
                   subset = ~ gopgovernor == 1, 
                   weight = ~ WEIGHT, 
                   cluster = ~ stateid, 
                   data = das_saho)


# compute squared residuals
das_G = das_gatherings %>%
  mutate(ressq = as_gatherings$residuals^2)
das_GR = das_gatherings %>%
  filter(gopgovernor == 1) %>%
  mutate(ressq = as_gatherings_R$residuals^2)
das_GD = das_gatherings %>%
  filter(demgovernor == 1) %>%
  mutate(ressq = as_gatherings_D$residuals^2)

das_S = das_saho %>%
  mutate(ressq = as_saho$residuals^2)
das_SR = das_saho %>%
  filter(gopgovernor == 1) %>%
  mutate(ressq = as_saho_R$residuals^2)
das_SD = das_saho %>%
  filter(demgovernor == 1) %>%
  mutate(ressq = as_saho_D$residuals^2)

# build weights

# Sum of suqared residuals
tg = sum(das_G$ressq)
tgD = sum(das_GD$ressq)
tgR = sum(das_GR$ressq)

ts = sum(das_S$ressq)
tsD = sum(das_SD$ressq)
tsR = sum(das_SR$ressq)

# Total N. Respondents
tg_resp = nrow(das_G)
tgD_resp = sum(das_GD$dem)
tgR_resp = sum(das_GR$gop)

ts_resp = nrow(das_S)
tsD_resp = sum(das_SD$dem)
tsR_resp = sum(das_SR$gop)


das_G = das_G %>%
  group_by(stateid, state) %>%
  summarise(n_resp_G = n(),
            wG = sum(ressq)/tg,
            wG_N = n()/tg_resp)

das_GR = das_GR %>%
  group_by(stateid, state) %>%
  summarise(n_resp_GR = n(),
            wGR = sum(ressq)/tgR,
            wGR_N = n()/tgR_resp)

das_GD = das_GD %>%
  group_by(stateid, state) %>%
  summarise(n_resp_GD = n(),
            wGD = sum(ressq)/tgD,
            wGD_N = n()/tgD_resp)

das_S = das_S %>%
  group_by(stateid, state) %>%
  summarise(n_resp_S = n(),
            wS = sum(ressq)/ts,
            wS_N = n()/ts_resp)

das_SR = das_SR %>%
  group_by(stateid, state) %>%
  summarise(n_resp_SR = n(),
            wSR = sum(ressq)/tsR,
            wSR_N = n()/tsR_resp)

das_SD = das_SD %>%
  group_by(stateid, state) %>%
  summarise(n_resp_SD = n(),
            wSD = sum(ressq)/tsD,
            wSD_N = n()/tsD_resp)

dastate = das_G %>%
  left_join(., das_GR) %>%
  left_join(., das_GD) %>%
  left_join(., das_S) %>%
  left_join(., das_SR) %>%
  left_join(., das_SD)


us_states <- get_acs(
  geography = "state",
  variables = "B15001_001E",
  year = 2021,
  survey = "acs1",
  geometry = TRUE,
  resolution = "20m"
) %>%
  shift_geometry()

# Merge spatial data with state-level data
df_states <- us_states %>%
  filter(!NAME %in% c("District of Columbia", "Puerto Rico")) %>%
  left_join(., dastate, by = c("NAME" = "state"))

df_stp = bind_rows(df_states %>% mutate(val = wS, name = "Effective Sample\nStay-at-Home\n"),
                   df_states %>% mutate(val = wG, name = "Effective Sample\nGatherings\n"),
                   df_states %>% mutate(val = wG_N, name = "Nominal Sample\nShare of Respondents by State\n"))

df_stp$name = factor(df_stp$name, levels = rev(unique(df_stp$name)))

ggplot(df_stp, fill = val) +
  geom_sf(aes(fill = val)) +
  scale_fill_distiller(palette = "Greens", 
                       direction = 1, 
                       breaks = c(0, .2, .4), 
                       limits = c(0, .4)) +
  theme_void() +
  facet_grid(~name) +
  theme(panel.background = element_blank(),
        axis.text = element_blank(),
        legend.position = "bottom",
        legend.title = element_text(vjust = .75, size = 7),
        legend.text = element_text(size = 6),
        legend.key.height= unit(.3, 'cm'),
        legend.key.width= unit(.75, 'cm'),
        legend.frame = element_rect(color = "black", fill = NA)
  ) +
  labs(fill = "Weight")

ggsave("figures/Figure S25.png", bg = "white", width = 6.5, height = 2.8)


